import cv2
import numpy as np
import os
from config import TEMPLATE_PATH, MATCH_THRESHOLD

def load_template(template_name):
    """
    加载模板图片
    
    Args:
        template_name (str): 模板图片文件名
    
    Returns:
        numpy.ndarray: 模板图片数据
    """
    template_path = os.path.join(TEMPLATE_PATH, template_name)
    if not os.path.exists(template_path):
        raise FileNotFoundError(f"模板图片未找到: {template_path}")
    
    template = cv2.imread(template_path)
    if template is None:
        raise ValueError(f"无法加载模板图片: {template_path}")
    
    return template

def match_template(screen_image, template_image, threshold=MATCH_THRESHOLD):
    """
    在屏幕截图中查找模板图片
    
    Args:
        screen_image (numpy.ndarray): 屏幕截图
        template_image (numpy.ndarray): 模板图片
        threshold (float): 匹配阈值
    
    Returns:
        tuple: (是否匹配成功, 匹配位置, 匹配置信度)
    """
    # 转换为灰度图
    screen_gray = cv2.cvtColor(screen_image, cv2.COLOR_BGR2GRAY)
    template_gray = cv2.cvtColor(template_image, cv2.COLOR_BGR2GRAY)
    
    # 执行模板匹配
    result = cv2.matchTemplate(screen_gray, template_gray, cv2.TM_CCOEFF_NORMED)
    
    # 查找最佳匹配位置
    min_val, max_val, min_loc, max_loc = cv2.minMaxLoc(result)
    
    # 判断是否匹配成功
    if max_val >= threshold:
        # 计算匹配区域的中心点
        center_x = max_loc[0] + template_gray.shape[1] // 2
        center_y = max_loc[1] + template_gray.shape[0] // 2
        return True, (center_x, center_y), max_val
    else:
        return False, None, max_val

def match_multiple_templates(screen_image, template_names, threshold=MATCH_THRESHOLD):
    """
    在屏幕截图中查找多个模板图片
    
    Args:
        screen_image (numpy.ndarray): 屏幕截图
        template_names (list): 模板图片文件名列表
        threshold (float): 匹配阈值
    
    Returns:
        list: 匹配结果列表 [(模板名, 是否匹配成功, 匹配位置, 匹配置信度), ...]
    """
    results = []
    
    for template_name in template_names:
        try:
            template = load_template(template_name)
            is_match, position, confidence = match_template(screen_image, template, threshold)
            results.append((template_name, is_match, position, confidence))
        except Exception as e:
            print(f"处理模板 {template_name} 时出错: {e}")
            results.append((template_name, False, None, 0))
    
    return results

def create_test_template():
    """
    创建一个简单的测试模板图片
    """
    # 创建一个简单的红色圆形作为测试模板
    template = np.zeros((100, 100, 3), dtype=np.uint8)
    cv2.circle(template, (50, 50), 30, (0, 0, 255), -1)  # 红色圆形
    cv2.imwrite(os.path.join(TEMPLATE_PATH, "test_template.png"), template)
    print("已创建测试模板图片: test_template.png")

def main():
    # 测试模板匹配功能
    # 注意：这需要实际的屏幕截图和模板图片才能运行
    print("模板匹配模块测试")
    print("请确保在templates目录中有模板图片，并提供屏幕截图进行测试")
    
    # 创建测试模板
    if not os.path.exists(TEMPLATE_PATH):
        os.makedirs(TEMPLATE_PATH)
    
    template_files = [f for f in os.listdir(TEMPLATE_PATH)
                     if f.lower().endswith(('.png', '.jpg', '.jpeg'))]
    
    if not template_files:
        print("未找到模板图片，创建测试模板...")
        create_test_template()

if __name__ == "__main__":
    main()